package com.smart.translator.util;

import android.util.Log;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import okio.ByteString;
import org.json.JSONException;
import org.json.JSONObject;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

public class AsrWsClient {
    private static final String TAG = "AsrWsClient";
    
    // Protocol constants - 与BigASRWebsocketDemo保持一致
    private static final byte PROTOCOL_VERSION = 0b0001;
    private static final byte DEFAULT_HEADER_SIZE = 0b0001;
    
    // Message Type
    private static final byte FULL_CLIENT_REQUEST = 0b0001;
    private static final byte AUDIO_ONLY_REQUEST = 0b0010;
    private static final byte FULL_SERVER_RESPONSE = 0b1001;
    private static final byte SERVER_ACK = 0b1011;
    private static final byte SERVER_ERROR_RESPONSE = 0b1111;
    
    // Message Type Specific Flags
    private static final byte NO_SEQUENCE = 0b0000;
    private static final byte POS_SEQUENCE = 0b0001;
    private static final byte NEG_SEQUENCE = 0b0010;
    private static final byte NEG_WITH_SEQUENCE = 0b0011;
    
    // Message Serialization
    private static final byte NO_SERIALIZATION = 0b0000;
    private static final byte JSON = 0b0001;
    
    // Message Compression
    private static final byte NO_COMPRESSION = 0b0000;
    private static final byte GZIP = 0b0001;
    
    // Configuration
    private String wsUrl = "wss://openspeech.bytedance.com/api/v3/sauc/bigmodel";

    //private String wsUrl = "wss://openspeech.bytedance.com/api/v3/sauc/bigmodel_nostream";


    private String appKey = "1290154717";
    private String accessKey = "XuNsfxpCTjct9Y1r7fpH4B9HmTk_LWf6";
    private String resourceId = "volc.bigasr.sauc.duration";
    
    // Audio configuration - 默认使用16kHz（重采样后的采样率）
    private String uid = "test";
    private String format = "pcm";
    private int rate = 16000; // 默认16kHz，支持8kHz输入自动重采样
    private int bits = 16;
    private int channel = 1;
    private String codec = "raw";
    
    // WebSocket
    private WebSocket webSocket;
    private boolean isConnected = false;
    private AsrWsCallback callback;
    private int seq = 1;
    private boolean initialRequestSent = false;
    private boolean hasError = false;
    private boolean waitingForAck = false; // 等待ACK状态
    
    // 音频数据缓冲
    private ByteArrayOutputStream audioBuffer = new ByteArrayOutputStream();
    private static final int CHUNK_SIZE = 3200; // 200ms的16kHz音频数据：16k * 16bit * 1声道 * 0.2秒 = 3200字节
    
    // 默认启用重采样
    private boolean enableResampling = true;
    private AudioResampler audioResampler;
    
    // 音频统计
    private long totalAudioBytes = 0;
    private int audioChunkCount = 0;
    
    public interface AsrWsCallback {
        void onConnected();
        void onInitialRequestSent();
        void onMessage(String text, boolean isFinal);
        void onError(String error);
        void onClosed();
    }
    
    public AsrWsClient() {
        // 默认启用重采样
        enableResampling();
    }
    
    public AsrWsClient(String wsUrl, String appKey, String accessKey, String resourceId) {
        this.wsUrl = wsUrl;
        this.appKey = appKey;
        this.accessKey = accessKey;
        this.resourceId = resourceId;
        // 默认启用重采样
        enableResampling();
    }
    
    public void setCallback(AsrWsCallback callback) {
        this.callback = callback;
    }
    
    /**
     * 启用音频重采样（8kHz -> 16kHz）
     */
    public void enableResampling() {
        this.enableResampling = true;
        this.audioResampler = new AudioResampler();
        this.rate = AudioResampler.SAMPLE_RATE_16K; // 更新采样率为16kHz
        Log.d(TAG, "Audio resampling enabled - 8kHz -> 16kHz");
    }
    
    /**
     * 禁用音频重采样
     */
    public void disableResampling() {
        this.enableResampling = false;
        this.audioResampler = null;
        this.rate = 8000; // 恢复为8kHz
        Log.d(TAG, "Audio resampling disabled");
    }
    
    /**
     * 设置音频配置
     */
    public void setAudioConfig(int sampleRate, int bits, int channels) {
        this.rate = sampleRate;
        this.bits = bits;
        this.channel = channels;
        Log.d(TAG, "Audio config updated - sampleRate: " + sampleRate + 
                   ", bits: " + bits + ", channels: " + channels);
    }
    
    public void setAudioConfig(int sampleRate) {
        setAudioConfig(sampleRate, 16, 1);
    }
    
    /**
     * 获取音频统计信息
     */
    public String getAudioStats() {
        return String.format("Total audio bytes: %d, Chunks sent: %d", totalAudioBytes, audioChunkCount);
    }
    
    public void connect() {
        if (isConnected) {
            Log.w(TAG, "WebSocket already connected");
            return;
        }
        
        // 重置状态
        resetAudioBuffer();
        
        OkHttpClient client = new OkHttpClient.Builder()
            .pingInterval(50, TimeUnit.SECONDS)
            .readTimeout(100, TimeUnit.SECONDS)
            .writeTimeout(100, TimeUnit.SECONDS)
            .build();
            
        String connectId = UUID.randomUUID().toString();
        
        Log.d(TAG, "Connecting to ASR with parameters:");
        Log.d(TAG, "  URL: " + wsUrl);
        Log.d(TAG, "  AppKey: " + appKey.substring(0, Math.min(8, appKey.length())) + "...");
        Log.d(TAG, "  AccessKey: " + accessKey.substring(0, Math.min(8, accessKey.length())) + "...");
        Log.d(TAG, "  ResourceId: " + resourceId);
        Log.d(TAG, "  ConnectId: " + connectId);
        Log.d(TAG, "  Sample Rate: " + rate + "Hz (Resampling: " + enableResampling + ")");
        
        Request request = new Request.Builder()
            .url(wsUrl)
            .header("X-Api-App-Key", appKey)
            .header("X-Api-Access-Key", accessKey)
            .header("X-Api-Resource-Id", resourceId)
            .header("X-Api-Connect-Id", connectId)
            .build();
            
        webSocket = client.newWebSocket(request, new WebSocketListener() {
            @Override
            public void onOpen(WebSocket webSocket, Response response) {
                Log.d(TAG, "ASR WebSocket connected");
                isConnected = true;
                if (callback != null) {
                    callback.onConnected();
                }
                
                // 不在这里自动发送初始请求，让外部控制
            }
            
            @Override
            public void onMessage(WebSocket webSocket, String text) {
                Log.d(TAG, "Received text message: " + text);
                // 处理文本消息（如心跳响应等）
                try {
                    JSONObject json = new JSONObject(text);
                    if (json.has("error")) {
                        String errorMsg = json.getString("error");
                        Log.e(TAG, "Server error: " + errorMsg);
                        if (callback != null) {
                            callback.onError("Server error: " + errorMsg);
                        }
                    }
                } catch (JSONException e) {
                    Log.d(TAG, "Non-JSON text message: " + text);
                }
            }
            
            @Override
            public void onMessage(WebSocket webSocket, ByteString bytes) {
                //Log.d(TAG, "Received binary message, length: " + bytes.size());
                try {
                    int sequence = parseResponse(bytes.toByteArray());
                    
                    if (sequence == -999) {
                        // 服务器错误，关闭连接
                        Log.d(TAG, "Server error received, closing connection");
                        webSocket.close(1000, "server error");
                        return;
                    }
                    
                    boolean isLastPackage = sequence < 0;
                    if (isLastPackage) {
                        Log.d(TAG, "Received last package, closing connection");
                        webSocket.close(1000, "finished");
                        return;
                    }
                    
                } catch (Exception e) {
                    Log.e(TAG, "Error parsing response", e);
                    if (callback != null) {
                        callback.onError("Error parsing response: " + e.getMessage());
                    }
                }
            }
            
            @Override
            public void onClosing(WebSocket webSocket, int code, String reason) {
                Log.d(TAG, "WebSocket closing: " + code + " " + reason);
            }
            
            @Override
            public void onClosed(WebSocket webSocket, int code, String reason) {
                Log.d(TAG, "WebSocket closed: " + code + " " + reason);
                isConnected = false;
                if (callback != null) {
                    callback.onClosed();
                }
            }
            
            @Override
            public void onFailure(WebSocket webSocket, Throwable t, Response response) {
                Log.e(TAG, "WebSocket failure", t);
                Log.e(TAG, "WebSocket failure details - Response: " + (response != null ? response.code() : "null") + 
                       ", Message: " + t.getMessage() + ", Cause: " + (t.getCause() != null ? t.getCause().getMessage() : "null"));
                isConnected = false;
                if (callback != null) {
                    callback.onError("WebSocket failure: " + t.getMessage() + " (Code: " + (response != null ? response.code() : "unknown") + ")");
                }
            }
        });
    }
    
    public void sendInitialRequest() {
        if (!isConnected || webSocket == null) {
            Log.e(TAG, "WebSocket not connected");
            return;
        }
        
        if (initialRequestSent) {
            Log.w(TAG, "Initial request already sent, skipping");
            return;
        }
        
        try {
            // 构建请求参数
            JSONObject user = new JSONObject();
            user.put("uid", uid);
            
            JSONObject audio = new JSONObject();
            audio.put("format", format);
            audio.put("sample_rate", rate); // 使用重采样后的采样率
            audio.put("bits", bits);
            audio.put("channel", channel);
            audio.put("codec", codec);
            
            JSONObject request = new JSONObject();
            request.put("model_name", "bigmodel");
            request.put("result_type", "single");
            request.put("enable_punc", true);
            
            JSONObject payload = new JSONObject();
            payload.put("user", user);
            payload.put("audio", audio);
            payload.put("request", request);
            
            String payloadStr = payload.toString();
            Log.d(TAG, "Initial request payload: " + payloadStr);
            
            // 压缩payload
            byte[] payloadBytes = gzipCompress(payloadStr.getBytes());
            
            // 组装fullClientRequest
            byte[] header = getHeader(FULL_CLIENT_REQUEST, POS_SEQUENCE, JSON, GZIP, (byte) 0);
            byte[] payloadSize = intToBytes(payloadBytes.length);
            byte[] seqBytes = generateBeforePayload(seq);
            
            byte[] fullClientRequest = new byte[header.length + seqBytes.length + payloadSize.length + payloadBytes.length];
            int destPos = 0;
            System.arraycopy(header, 0, fullClientRequest, destPos, header.length);
            destPos += header.length;
            System.arraycopy(seqBytes, 0, fullClientRequest, destPos, seqBytes.length);
            destPos += seqBytes.length;
            System.arraycopy(payloadSize, 0, fullClientRequest, destPos, payloadSize.length);
            destPos += payloadSize.length;
            System.arraycopy(payloadBytes, 0, fullClientRequest, destPos, payloadBytes.length);
            
            boolean success = webSocket.send(ByteString.of(fullClientRequest));
            if (success) {
                Log.d(TAG, "Sent initial request successfully");
                initialRequestSent = true;
                waitingForAck = true;
                if (callback != null) {
                    callback.onInitialRequestSent();
                }
            } else {
                Log.e(TAG, "Failed to send initial request");
            }
            
        } catch (Exception e) {
            Log.e(TAG, "Error sending initial request", e);
            if (callback != null) {
                callback.onError("Error sending initial request: " + e.getMessage());
            }
        }
    }
    
    public void sendAudioData(byte[] audioData, boolean isLast) {
        if (!canSendAudioData()) {
            Log.w(TAG, "Cannot send audio data - connected: " + isConnected + 
                       ", hasError: " + hasError + 
                       ", initialRequestSent: " + initialRequestSent + 
                       ", waitingForAck: " + waitingForAck);
            return;
        }
        
        try {
            byte[] processedAudioData = audioData;
            
            // 如果启用了重采样，进行8kHz -> 16kHz转换
            if (enableResampling && audioResampler != null) {
                processedAudioData = audioResampler.resample8kTo16k(audioData);
                //Log.d(TAG, "Audio resampled: " + audioData.length + " bytes -> " + processedAudioData.length + " bytes");
            }
            
            // 更新统计信息
            totalAudioBytes += processedAudioData.length;
            
            // 添加到缓冲区
            audioBuffer.write(processedAudioData);
            
            // 按CHUNK_SIZE分包发送
            while (audioBuffer.size() >= CHUNK_SIZE) {
                byte[] allData = audioBuffer.toByteArray();
                byte[] chunk = Arrays.copyOfRange(allData, 0, CHUNK_SIZE);
                
                // 移除已发送的数据
                audioBuffer.reset();
                audioBuffer.write(allData, CHUNK_SIZE, allData.length - CHUNK_SIZE);
                
                // 发送音频块
                sendAudioChunk(chunk, false);
                audioChunkCount++;
                
                // 音频帧不需要等待ACK，可以连续发送
            }
            
            // 如果是最后一帧且有剩余数据，发送剩余数据
            if (isLast && audioBuffer.size() > 0) {
                byte[] remainingData = audioBuffer.toByteArray();
                audioBuffer.reset();
                sendAudioChunk(remainingData, true);
                audioChunkCount++;
            } else if (isLast) {
                // 没有剩余数据，发送空结束帧
                sendAudioChunk(new byte[0], true);
                audioChunkCount++;
            }
            
            //Log.d(TAG, "Audio stats: " + getAudioStats());
            
        } catch (Exception e) {
            Log.e(TAG, "Error sending audio data", e);
            if (callback != null) {
                callback.onError("Error sending audio data: " + e.getMessage());
            }
        }
    }
    
    private void sendAudioChunk(byte[] audioData, boolean isLast) {
        if (webSocket == null || !isConnected) {
            Log.e(TAG, "WebSocket not connected, cannot send audio chunk");
            return;
        }
        
        seq++;
        // Log.d(TAG, "Sending audio chunk, seq: " + seq + ", isLast: " + isLast + 
        //            ", data size: " + audioData.length + " bytes");
        
        // 修正序列号处理逻辑，与官方demo保持一致
        int finalSeq = isLast ? -seq : seq;
        
        byte messageTypeSpecificFlags = isLast ? NEG_WITH_SEQUENCE : POS_SEQUENCE;
        
        // header
        byte[] header = getHeader(AUDIO_ONLY_REQUEST, messageTypeSpecificFlags, JSON, GZIP, (byte) 0);
        
        // sequence
        byte[] sequenceBytes = generateBeforePayload(finalSeq);
        
        // payload
        byte[] payloadBytes = gzipCompress(audioData, audioData.length);
        byte[] payloadSize = intToBytes(payloadBytes.length);
        
        byte[] audioOnlyRequest = new byte[header.length + sequenceBytes.length + payloadSize.length + payloadBytes.length];
        int destPos = 0;
        System.arraycopy(header, 0, audioOnlyRequest, destPos, header.length);
        destPos += header.length;
        System.arraycopy(sequenceBytes, 0, audioOnlyRequest, destPos, sequenceBytes.length);
        destPos += sequenceBytes.length;
        System.arraycopy(payloadSize, 0, audioOnlyRequest, destPos, payloadSize.length);
        destPos += payloadSize.length;
        System.arraycopy(payloadBytes, 0, audioOnlyRequest, destPos, payloadBytes.length);
        
        try {
            boolean success = webSocket.send(ByteString.of(audioOnlyRequest));
            if (!success) {
                Log.e(TAG, "Failed to send audio chunk");
                hasError = true;
                if (callback != null) {
                    callback.onError("Failed to send audio chunk");
                }
            }
        } catch (Exception e) {
            Log.e(TAG, "Exception sending audio chunk", e);
            hasError = true;
            if (callback != null) {
                callback.onError("Exception sending audio chunk: " + e.getMessage());
            }
        }
    }
    
    public void resetAudioBuffer() {
        try {
            if (audioBuffer != null) {
                audioBuffer.reset();
            }
        } catch (Exception e) {
            Log.e(TAG, "重置音频缓冲区异常", e);
        }
        seq = 1;
        initialRequestSent = false;
        hasError = false;
        waitingForAck = false;
        
        // 重置统计信息
        totalAudioBytes = 0;
        audioChunkCount = 0;
        
        // 重置重采样器
        if (audioResampler != null) {
            audioResampler.clearBuffer();
        }
    }
    
    public boolean isConnected() {
        return isConnected && webSocket != null;
    }
    
    public boolean canSendAudioData() {
        return isConnected && !hasError && initialRequestSent;
    }
    
    public void close() {
        // 清理音频缓冲区
        try {
            if (audioBuffer != null) {
                audioBuffer.close();
            }
        } catch (Exception e) {
            Log.e(TAG, "清理音频缓冲区异常", e);
        }
        
        if (webSocket != null) {
            webSocket.close(1000, null);
            webSocket = null;
        }
        isConnected = false;
    }
    
    // 协议相关方法 - 与BigASRWebsocketDemo保持一致
    private static byte[] getHeader(byte messageType, byte messageTypeSpecificFlags, byte serialMethod, byte compressionType, byte reservedData) {
        final byte[] header = new byte[4];
        header[0] = (byte) ((PROTOCOL_VERSION << 4) | DEFAULT_HEADER_SIZE);
        header[1] = (byte) ((messageType << 4) | messageTypeSpecificFlags);
        header[2] = (byte) ((serialMethod << 4) | compressionType);
        header[3] = reservedData;
        return header;
    }
    
    private static byte[] intToBytes(int a) {
        return new byte[]{
            (byte) ((a >> 24) & 0xFF),
            (byte) ((a >> 16) & 0xFF),
            (byte) ((a >> 8) & 0xFF),
            (byte) (a & 0xFF)
        };
    }
    
    private static int bytesToInt(byte[] src) {
        if (src == null || src.length != 4) {
            throw new IllegalArgumentException("Invalid byte array length");
        }
        return ((src[0] & 0xFF) << 24)
                | ((src[1] & 0xff) << 16)
                | ((src[2] & 0xff) << 8)
                | ((src[3] & 0xff));
    }
    
    private static byte[] generateBeforePayload(int seq) {
        return intToBytes(seq);
    }
    
    private static byte[] gzipCompress(byte[] src) {
        return gzipCompress(src, src.length);
    }
    
    private static byte[] gzipCompress(byte[] src, int len) {
        if (src == null || len == 0) {
            return new byte[0];
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        GZIPOutputStream gzip = null;
        try {
            gzip = new GZIPOutputStream(out);
            gzip.write(src, 0, len);
        } catch (IOException e) {
            Log.e(TAG, "GZIP compression error", e);
        } finally {
            if (gzip != null) {
                try {
                    gzip.close();
                } catch (IOException e) {
                    Log.e(TAG, "Error closing GZIP stream", e);
                }
            }
        }
        return out.toByteArray();
    }
    
    private static byte[] gzipDecompress(byte[] src) {
        if (src == null || src.length == 0) {
            return null;
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ByteArrayInputStream ins = new ByteArrayInputStream(src);
        GZIPInputStream gzip = null;
        try {
            gzip = new GZIPInputStream(ins);
            byte[] buffer = new byte[1024];
            int len;
            while ((len = gzip.read(buffer)) > 0) {
                out.write(buffer, 0, len);
            }
        } catch (IOException e) {
            Log.e(TAG, "GZIP decompression error", e);
        } finally {
            if (gzip != null) {
                try {
                    gzip.close();
                } catch (IOException e) {
                    Log.e(TAG, "Error closing GZIP stream", e);
                }
            }
        }
        return out.toByteArray();
    }
    
    private int parseResponse(byte[] res) {
        if (res == null || res.length == 0) {
            return -1;
        }
        
        final byte num = 0b00001111;
        
        // header 32 bit=4 byte
        int protocolVersion = (res[0] >> 4) & num;
        int headerSize = res[0] & 0x0f;
        int messageType = (res[1] >> 4) & num;
        int messageTypeSpecificFlags = res[1] & 0x0f;
        int serializationMethod = (res[2] >> 4) & num;
        int messageCompression = res[2] & 0x0f;
        int reserved = res[3];
        
        // Log.d(TAG, "Parsed header - protocolVersion: " + protocolVersion + 
        //            ", headerSize: " + headerSize + 
        //            ", messageType: " + messageType + 
        //            ", messageTypeSpecificFlags: " + messageTypeSpecificFlags + 
        //            ", serializationMethod: " + serializationMethod + 
        //            ", messageCompression: " + messageCompression);
        
        // sequence 4 byte
        byte[] temp = new byte[4];
        System.arraycopy(res, 4, temp, 0, temp.length);
        int sequence = bytesToInt(temp);
        
        // payload size 4 byte
        System.arraycopy(res, 8, temp, 0, temp.length);
        int payloadSize = bytesToInt(temp);
        
        byte[] payload = new byte[res.length - 12];
        System.arraycopy(res, 12, payload, 0, payload.length);
        
        String payloadStr = null;
        
        // 根据消息类型处理
        if (messageType == FULL_SERVER_RESPONSE) {
            if (messageCompression == GZIP) {
                payloadStr = new String(gzipDecompress(payload));
            } else {
                payloadStr = new String(payload);
            }
            //Log.d(TAG, "FULL_SERVER_RESPONSE payload: " + payloadStr);
            
            // 收到初始响应，可以开始发送音频数据
            waitingForAck = false;
            
            // 解析识别结果
            try {
                JSONObject json = new JSONObject(payloadStr);
                String text = extractTextFromResponse(json);
                if (text != null && !text.isEmpty()) {
                    boolean isFinal = json.optBoolean("is_final", false);
                    Log.d(TAG, "ASR recognition result - text: '" + text + "', isFinal: " + isFinal);
                    if (callback != null) {
                        callback.onMessage(text, isFinal);
                    }
                }
            } catch (JSONException e) {
                Log.e(TAG, "Error parsing FULL_SERVER_RESPONSE JSON", e);
            }
            
        } else if (messageType == SERVER_ACK) {
            payloadStr = new String(payload);
            Log.d(TAG, "SERVER_ACK payload: " + payloadStr);
            
            // 收到ACK，重置等待状态（主要用于初始请求）
            waitingForAck = false;
            
        } else if (messageType == SERVER_ERROR_RESPONSE) {
            // 此时 sequence 含义就是 错误码 code，payload 就是 error msg
            payloadStr = new String(payload);
            Log.e(TAG, "SERVER_ERROR_RESPONSE - code: " + sequence + ", error: " + payloadStr);
            hasError = true;
            waitingForAck = false;
            if (callback != null) {
                callback.onError("Server error: code=" + sequence + ", message=" + payloadStr);
            }
            // 返回特殊值表示错误，不继续发送音频数据
            return -999;
        }
        
        return sequence;
    }
    
    private String extractTextFromResponse(JSONObject json) {
        try {
            // 豆包API的响应格式可能是以下几种：
            // 1. {"result": {"text": "识别文本"}}
            // 2. {"result": [{"text": "文本1"}, {"text": "文本2"}]}
            // 3. {"text": "识别文本"}
            
            if (json.has("result")) {
                Object result = json.get("result");
                if (result instanceof JSONObject) {
                    // 格式1: {"result": {"text": "识别文本"}}
                    JSONObject resultObj = (JSONObject) result;
                    if (resultObj.has("text")) {
                        return resultObj.getString("text");
                    }
                } else if (result instanceof org.json.JSONArray) {
                    // 格式2: {"result": [{"text": "文本1"}, {"text": "文本2"}]}
                    org.json.JSONArray resultArray = (org.json.JSONArray) result;
                    StringBuilder textBuilder = new StringBuilder();
                    for (int i = 0; i < resultArray.length(); i++) {
                        JSONObject segment = resultArray.getJSONObject(i);
                        if (segment.has("text")) {
                            textBuilder.append(segment.getString("text"));
                        }
                    }
                    return textBuilder.toString();
                }
            } else if (json.has("text")) {
                // 格式3: {"text": "识别文本"}
                return json.getString("text");
            }
        } catch (Exception e) {
            Log.e(TAG, "Error extracting text from response", e);
        }
        return null;
    }
} 